import torch
import torch.nn as nn

class SegModel(nn.Module):
    def __init__(self, plm_model, plm_config, vocab):
        super(SegModel, self).__init__()
        self.plm_model = plm_model

        self.output_layer = nn.Linear(
            in_features=plm_config.hidden_size, 
            out_features=vocab.label_size
        )

    def forward(self, inputs):
        outputs = self.plm_model(**inputs)
        logits = self.output_layer(outputs.last_hidden_state)

        return logits